//
//  fooplugin.cpp
//  fix_lldb_test
//
//  Created by marcow on 2021/11/12.
//

#include <LLDB/LLDB.h>
#include <dlfcn.h>
#include <mach-o/dyld.h>
#include <mach-o/nlist.h>


using namespace lldb;

static void dsprintf(const char *format, ...) {
  static int quiet_mode = 0;
  if (quiet_mode == 0) {
    quiet_mode = getenv("DS_QUIET") ? 1 : -1;
  }
  if (quiet_mode == 1) { return; }
  printf("\033[36;1m");
  va_list args;
  va_start(args, format);
  vprintf(format, args);
  va_end( args );
  printf("\033[0m");
}

extern "C" {
void* xcselect_get_developer_dir_path(char *ptr, size_t length, char *a, char*b, char *c);
}

// expects code to be linked to /usr/lib/xcselect.dylib ("xcode-select -p" equivalent)
static char *xcode_path() {
  char *input = (char *)calloc(0x400, sizeof(char));
  char dunno1, dunno2, dunno3;

  xcselect_get_developer_dir_path(input, 0x400, &dunno1, &dunno2, &dunno3);
  return input;
}

/// Searches in LLDB.framework for symbols
/// Doesn't use dlopen/dlsym since these functions are non exported :[
static uintptr_t address_for_function(char *symbol_name) {
  static void *handle = NULL;
  static int LLDBIndex = 0;
  static struct mach_header_64* header;
  static struct symtab_command *symtab_cmd = NULL;
  static struct segment_command_64 *linkedit_cmd = NULL;
  static intptr_t slide;
  static long onceToken;
  static const char *fullpathToLLDB = NULL;
  /////////////////////////////////
  
  // a poor man's dispatch_once :]
  if (onceToken == 0) {
    onceToken = -1;
    
    // First need to determine the path of Xcode, don't assume it's always /Applications/Xcode.app/
    char path[0x400] = {};
    strcpy(path, xcode_path());
    // This will give me something like /Applications/Xcode.app/Contents/Developer
    // replace the ./Developer with the ./SharedFrameworks/LLDB...
    char *offset = strstr(path, "/Developer");
    if (!offset) {
      dsprintf("unable to parse \"%s\"\n", path);
      return 0;
    }
    strcpy(offset, "/SharedFrameworks/LLDB.framework/Versions/A/LLDB");
    handle = dlopen(path, RTLD_NOW);
    if (!handle) {
      dsprintf("unable find LLDB.framework at \"%s\"\n", path);
      return 0;
    }
    
    // Find the LLDB.framework loaded into the LLDB process
    for (int i = 0; i < _dyld_image_count(); i++) {
      const char *imagePath = _dyld_get_image_name(i);
      if (strstr(imagePath, "SharedFrameworks/LLDB.framework/Versions/A/LLDB") != NULL) {
        LLDBIndex = i;
        fullpathToLLDB = imagePath;
        break;
      }
    }
    
    if (!fullpathToLLDB) {
      dsprintf("unable find path to LLDB.framework\n");
      return 0;
    }
    
    // Get the Mach-O header's memory address of LLDB.framework
    header = (struct mach_header_64*)_dyld_get_image_header(LLDBIndex);
    if (header->magic != MH_MAGIC_64) {
      dsprintf("LLDB.framework header messed up\n");
      return 0;
    }
    
    slide = _dyld_get_image_vmaddr_slide(LLDBIndex);
    uintptr_t cur_pointer = (uintptr_t)header + sizeof(struct mach_header_64);
    
    // Get the LC_SYMTAB Load Command and linkedit section Load Command
    for (int i = 0; i < header->ncmds; i++) {
      struct load_command *cur_cmd =  (struct load_command *)cur_pointer;
      
      // LC_SYMTAB
      if (cur_cmd->cmd == LC_SYMTAB) {
        symtab_cmd = (struct symtab_command *)cur_cmd;
      }
      
      // LC_SEGMENT_64
      if (cur_cmd->cmd == LC_SEGMENT_64) {
        struct segment_command_64 *segment_cmd = (struct segment_command_64*)cur_cmd;
        if (strcmp(segment_cmd->segname, SEG_LINKEDIT) == 0) {
          linkedit_cmd = segment_cmd;
        }
      }
      
      cur_pointer += cur_cmd->cmdsize;
    }
    
    // Make sure we have the address to both of those Load Commands
    if (!symtab_cmd || !linkedit_cmd) {
      dsprintf("unable find the appropriate commands to find variables\n");
      return 0;
    }
  } // </poor man's dispatch_once>
  
  
  // If we got here, we found the location of the symbol table, search for the function
  uintptr_t base_pointer =  slide + linkedit_cmd->vmaddr - linkedit_cmd->fileoff;
  struct nlist_64* nlists = (struct nlist_64 *)(base_pointer + symtab_cmd->symoff);
  char* strtab = (char*)(base_pointer + symtab_cmd->stroff);
  int numSymbols = symtab_cmd->nsyms;
  
  for (int i = 0; i < numSymbols; i++) {
    struct nlist_64 symbol = nlists[i];
    char *str = &strtab[symbol.n_un.n_strx];
    if (strcmp(str, symbol_name) == 0) {
      uintptr_t resolved_address = (((uintptr_t)header) + symbol.n_value);
      return resolved_address;
    }
  }
  dsprintf("unable to find symbol \"%s\"\n", symbol_name);
  return 0;
}

namespace lldb {
bool PluginInitialize(SBDebugger debugger);
}

class ChildCommand : public SBCommandPluginInterface {
public:
    virtual bool DoExecute(SBDebugger debugger, char **command,
                           SBCommandReturnObject &result) {
        if (command) {
            const char *arg = *command;
            while (arg) {
                result.Printf("%s\n", arg);
                arg = *(++command);
            }
            return true;
        }
        return false;
    }
};

bool lldb::PluginInitialize(SBDebugger debugger) {
    lldb::SBCommandInterpreter interpreter = debugger.GetCommandInterpreter();
    lldb::SBCommand foo = interpreter.AddMultiwordCommand("foo", NULL);
    foo.AddCommand("child", new ChildCommand(), "a child of foo");

    SBCommandReturnObject result;
    const char *command = "version";
    interpreter.HandleCommand(command, result);
    if (result.IsValid()) {
        dsprintf("lldb version: %s\n", result.GetOutput());
    }
    result.Clear();
    
    const char *path = xcode_path();
    dsprintf("xcode path: %s\n", path);
    
    return true;
}

